import os
from typing import Optional, List
from elasticsearch import AsyncElasticsearch
from kg_retrieval.env_config import ES_CONFIG

# ES连接默认配置
DEFAULT_ES_HOSTS = ['https://localhost:9200']
DEFAULT_ES_USER = 'elastic'
DEFAULT_ES_PASS = ''
DEFAULT_ES_INDEX = 'jd_file_index_formal'

# 全局ES客户端实例
_es_client = None

async def get_es_client() -> Optional[AsyncElasticsearch]:
    """获取或创建ES客户端"""
    global _es_client
    
    # 如果ES未启用，返回None
    if not ES_CONFIG["ES_ENABLED"]:
        print("Elasticsearch功能未启用")
        return None
    
    if _es_client is None:
        try:
            # 使用环境变量配置
            es_hosts = ES_CONFIG["ES_HOSTS"]
            es_user = ES_CONFIG["ES_USER"]
            es_pass = ES_CONFIG["ES_PASS"]
            verify_certs = ES_CONFIG["ES_VERIFY_CERTS"]
            
            # 创建ES客户端
            _es_client = AsyncElasticsearch(
                hosts=es_hosts,
                basic_auth=(es_user, es_pass),
                verify_certs=verify_certs
            )
            
            # 测试连接
            await _es_client.info()
            print(f"Elasticsearch连接成功: {es_hosts}")
            
        except Exception as e:
            print(f"Elasticsearch连接失败: {e}")
            _es_client = None
    
    return _es_client

async def close_es_client():
    """关闭ES客户端连接"""
    global _es_client
    
    if _es_client:
        try:
            await _es_client.close()
            print("Elasticsearch连接已关闭")
        except Exception as e:
            print(f"关闭Elasticsearch连接出错: {e}")
        finally:
            _es_client = None

def get_es_index_name() -> str:
    """获取ES索引名称"""
    return ES_CONFIG["ES_INDEX"] 